package com.wuyanzu.shixicoderedo.common.threadpool;

import com.wuyanzu.shixicoderedo.common.TaskAction;
import com.wuyanzu.shixicoderedo.common.threadlocal.RequestContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 线程池管理器
 * 功能：
 * 1. 提供线程池执行任务
 * 2. 支持上下文传递
 * 3. 支持自定义拒绝策略
 */
public class ExecutorManagerWithContext {

    // 日志对象
    private static final Logger LOG = LoggerFactory.getLogger(ExecutorManagerWithContext.class);

    // CPU 核心数
    private static final int CPU = Math.max(1, Runtime.getRuntime().availableProcessors());

    // 创建线程池
    private static final ThreadPoolExecutor exec = new ThreadPoolExecutor(
            Math.max(2, CPU * 2),       // corePoolSize
            Math.max(4, CPU * 10),      // maximumPoolSize
            60L, TimeUnit.SECONDS,      // 空闲线程存活时间
            new LinkedBlockingQueue<>(200), // 队列容量
            new NamedThreadFactory("impr-task-%d"), // 自定义线程工厂
            new CustomRejectionPolicy(3, 200) // 自定义拒绝策略
    );

    // 单例实例
    private static final ExecutorManagerWithContext INSTANCE = new ExecutorManagerWithContext();

    // 私有构造，防止外部 new
    private ExecutorManagerWithContext() {}

    // 获取单例
    public static ExecutorManagerWithContext getInstance() { return INSTANCE; }

    /**
     * 自定义拒绝策略
     * 当线程池队列满时，尝试重试入队
     */
    public static class CustomRejectionPolicy implements RejectedExecutionHandler {

        // 日志
        private final Logger logger = LoggerFactory.getLogger(CustomRejectionPolicy.class);

        // 最大重试次数
        private final int maxRetries;

        // 每次重试等待时间（毫秒）
        private final long offerTimeoutMs;

        public CustomRejectionPolicy(int maxRetries, long offerTimeoutMs) {
            this.maxRetries = Math.max(0, maxRetries);
            this.offerTimeoutMs = Math.max(0, offerTimeoutMs);
        }

        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
            logger.warn("任务被拒绝，尝试重试 {} 次: {}", maxRetries, r);

            // 如果线程池关闭，直接丢弃
            if (executor.isShutdown() || executor.isTerminating()) {
                logger.error("线程池正在关闭，丢弃任务 {}", r);
                return;
            }

            // 获取队列引用
            BlockingQueue<Runnable> queue = executor.getQueue();

            int attempts = 0;

            // 循环重试
            while (attempts < maxRetries) {
                attempts++;
                try {
                    boolean offered = queue.offer(r, offerTimeoutMs, TimeUnit.MILLISECONDS);
                    if (offered) {
                        logger.info("第 {} 次重试成功，将任务重新入队: {}", attempts, r);
                        return;
                    } else {
                        logger.warn("第 {} 次重试失败", attempts);
                    }
                } catch (InterruptedException ie) {
                    Thread.currentThread().interrupt();
                    logger.warn("重试入队被中断", ie);
                    break;
                }
            }

            // 所有重试失败
            logger.error("所有 {} 次重试失败，任务丢弃: {}", attempts, r);
        }
    }

    /**
     * 命名线程工厂
     */
    static class NamedThreadFactory implements ThreadFactory {
        private final String format;
        private final AtomicInteger idx = new AtomicInteger(0);

        NamedThreadFactory(String format) { this.format = format; }

        @Override
        public Thread newThread(Runnable r) {
            Thread t = new Thread(r);
            t.setName(String.format(format, idx.getAndIncrement())); // 命名线程
            t.setDaemon(true); // 守护线程
            return t;
        }
    }

    /**
     * 并发执行一组 TaskAction
     *
     * 说明：
     * - 使用 CountDownLatch 等待所有任务完成（避免逐个 Future.get() 导致串行等待）
     * - 每个任务都用 ContextAwareCallable 包装，上下文（RequestContext）会注入到工作线程并在 finally 中清理
     * - 为保证 latch 能在任务抛异常时仍然减计数，实际提交的是一个在 finally 中调用 latch.countDown() 的 Callable
     *
     * @param ctx   请求上下文（会注入到每个执行线程的 ThreadLocal + MDC）
     * @param tasks 要并发执行的 TaskAction 列表
     * @param <T>   返回类型
     * @return 非空结果列表（按提交顺序的过滤版）
     * @throws InterruptedException 如果当前线程在等待时被中断
     * @throws ExecutionException   如果某个任务抛出异常则抛出 ExecutionException（包含 cause）
     */
    @SuppressWarnings("unchecked")
    public <T> List<T> executeTasks(RequestContext ctx,
                                    TaskAction<? extends T>... tasks) throws InterruptedException, ExecutionException {

        // 空任务直接返回不可变空列表
        if (tasks == null || tasks.length == 0) return List.of();

        // 为所有任务创建 CountDownLatch，保证 latch.await() 直到每个任务都执行到 finally
        final CountDownLatch latch = new CountDownLatch(tasks.length);

        // 保存 Future，以便在 latch.await() 之后统一收集结果
        List<Future<T>> futures = new ArrayList<>(tasks.length);

        // 保存最终非空结果
        List<T> results = new ArrayList<>(tasks.length);

        // 遍历任务并提交
        for (TaskAction<? extends T> raw : tasks) {
            // 将 TaskAction 转为 Callable（执行 doInAction）
            Callable<T> businessCallable = () -> ((TaskAction<T>) raw).doInAction();

            // 使用 ContextAwareCallable 包装，以便在子线程中设置 RequestContext + MDC，并在 finally 清理
            ContextAwareCallable<T> ctxCallable = new ContextAwareCallable<>(businessCallable, ctx);

            // 进一步包装一个 Callable，确保在任务结束（无论正常/异常/中断）都会执行 latch.countDown()
            Callable<T> latchAwareCallable = () -> {
                try {
                    return ctxCallable.call();
                } finally {
                    // 确保减计数，不会因为异常阻塞主线程等待
                    latch.countDown();
                }
            };

            // 提交到线程池
            Future<T> future = exec.submit(latchAwareCallable);

            // 保存 Future，稍后收集
            futures.add(future);
        }

        // 等待所有任务执行到 finally（即完成或抛异常或被取消）
        try {
            latch.await();
        } catch (InterruptedException ie) {
            // 恢复中断状态并向上抛出（调用方负责处理）
            Thread.currentThread().interrupt();
            throw ie;
        }

        // 所有任务已经完成或至少调用了 latch.countDown()，现在收集每个 Future 的结果
        for (Future<T> f : futures) {
            try {
                // Future.get() 此处不应该阻塞（任务已完成），但会抛出 ExecutionException 如果任务内部抛异常
                T r = f.get();
                if (r != null) results.add(r);
            } catch (InterruptedException ie) {
                // 在收集结果过程中被中断，恢复中断并抛出
                Thread.currentThread().interrupt();
                throw ie;
            } catch (ExecutionException ee) {
                // 将第一个遇到的执行异常向调用者抛出，包含原始异常信息
                throw ee;
            }
        }

        return results;
    }

    // 平滑关闭线程池
    public void shutdown() { exec.shutdown(); }

    // 立即关闭线程池
    public List<Runnable> shutdownNow() { return exec.shutdownNow(); }
}
